#include <vector>
#include <iostream>
#include <thread>
#include "pmath.h"

// cpu

// 生成随机数据
typedef struct {
    int start;
    int end;
    int seed;
    float *arr;
}Mes;

void rand_gen(Mes *mes) {
    
    srand((int)time(NULL) + mes->start + mes->seed);
    
    for (int i = mes->start; i < mes->end; i++) {
        int r = rand();
        float f = (float)r / RAND_MAX;
        float f2 = f * 2.0f - 1.0f;
        mes->arr[i] = f2;
    }
    
}

float * rand_arr(int size, int seed) {
    int thread_num = std::thread::hardware_concurrency();
    float *arr = new float[size];
    std::vector<std::thread> threads(thread_num);
    std::vector<Mes> mes(thread_num);
    int blocksdim;
    if (size%thread_num != 0) {
        if (size < thread_num) {
            for (int i = 0; i < size; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = Mes{start, end, seed, arr};
            }
        }else{
            blocksdim = size/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = size;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = Mes{start, end, seed, arr};
            }
        }
    }else{
        blocksdim = size/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = Mes{start, end, seed, arr};
            }
    }
    if (size < thread_num) {
        thread_num = size;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(rand_gen, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }

    return arr;
}

// 向量计算
typedef struct{
    int start;
    int end;
    int k;
    float *a;
    float *b;
    float *c;

}Pot;


void add_t(Pot * pot){
    float *a = pot->a;
    float *b = pot->b;
    float *c = pot->c;
    
    for(int i = pot->start; i < pot->end; i++){
        if (pot ->k == NULL){
            c[i] = a[i] + b[i];
        }else{
            c[i] = a[i] - b[i];
        }
        
    }
}

void mul_t(Pot *p){
    
    int start = p->start;
    int end = p->end;
    int k = p->k;
    float *A = p->a;
    float *B = p->b;
    float *C = p->c;

    for(int i = start; i < end; i++) {
        int row = i / k;
        int col = i % k;
        for (int j = 0; j < k; j++) {
            C[i] += A[row * k + j] * B[j * k + col];
        }
    }

}

// 加法中利用不到k值，将k值用于区分做加法还是减法
void add(float *a, float *b, float *c, int size, int tag){
    int thread_num = std::thread::hardware_concurrency();
    std::vector<std::thread> threads(thread_num);
    std::vector<Pot> mes(thread_num);
    int blocksdim;
    if (size%thread_num != 0) {
        if (size < thread_num) {
            for (int i = 0; i < size; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = Pot{start, end, tag, a, b, c };
            }
        }else{
            blocksdim = size/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = size;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = Pot{start, end, tag, a, b, c };
            }
        }
    }else{
        blocksdim = size/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = Pot{start, end, tag, a, b, c };
            }
    }
    if (size < thread_num) {
        thread_num = size;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(add_t, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
}
// 矩阵数乘
void lmul_t(Pot *p){
    float apha = *p->a;
    for (int i = p->start; i < p->end; i++){
        p->c[i] = apha * p->b[i];
    }
}

void lmul(float apha, float *b, float *c, int size){
    int thread_num = std::thread::hardware_concurrency();
    std::vector<std::thread> threads(thread_num);
    std::vector<Pot> mes(thread_num);
    int blocksdim;
    if (size%thread_num != 0) {
        if (size < thread_num) {
            for (int i = 0; i < size; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = Pot{start, end, NULL, &apha, b, c};
            }
        }else{
            blocksdim = size/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = size;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = Pot{start, end, NULL, &apha, b, c};
            }
        }
    }else{
        blocksdim = size/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = Pot{start, end, NULL, &apha, b, c};
            }
    }
    if (size < thread_num) {
        thread_num = size;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(lmul_t, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
}

// 矩阵乘法
void degmm(int m, int n, int k, float *a, float *b, float *c){
    int nums = m * n; // 矩阵元素个数
    int thread_num = std::thread::hardware_concurrency(); // 线程数量
    std::vector<std::thread> threads(thread_num);

    std::vector<Pot> pots(thread_num);

    int blocksdim;
    if (nums%thread_num != 0) {
        if (nums < thread_num) {
            for (int i = 0; i < nums; i++) {
                int end = i + 1;
                int start = i;
                pots[i] = Pot{start, end, k, a, b, c};
            }
        }else{
            blocksdim = nums/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = nums;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                pots[i] = Pot{start, end, k, a, b, c};
            }
        }
    }else{
        blocksdim = nums/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                pots[i] = Pot{start, end, k, a, b, c};
            }
    }
    if (nums < thread_num) {
        thread_num = nums;
    }

    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(mul_t, &pots[i]);
    }

    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
}
// 矩阵转置

typedef struct{
    int start;
    int end;
    int m;
    int n; 
    float *a;
    float *b;
}Pot_t;

void transpose_t(Pot_t *p){
    for(int idx=p->start;idx<p->end;idx++){
        int y=idx/p->n;
        int x=idx%p->n;
        p->b[x*p->m+y]=p->a[idx];
    }
}
void transpose(float *a, float *b, int m, int n){
    int size = m*n;
    int thread_num = std::thread::hardware_concurrency();
    std::vector<std::thread> threads(thread_num);
    std::vector<Pot_t> mes(thread_num);
    int blocksdim;
    if (size%thread_num != 0) {
        if (size < thread_num) {
            for (int i = 0; i < size; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = Pot_t{start, end, m, n, a, b};
            }
        }else{
            blocksdim = size/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = size;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = Pot_t{start, end, m, n, a, b};
            }
        }
    }else{
        blocksdim = size/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = Pot_t{start, end, m, n, a, b};
            }
    }
    if (size < thread_num) {
        thread_num = size;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(transpose_t, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
}

// abs

typedef struct{
    int start;
    int end;
    float *a;
    float *b;
}As;

void abs_t(As *p){
    for (int idx=p->start;idx<p->end;idx++) {
        if (p->a[idx] < 0) {
            p->b[idx] = -p->a[idx];
        }else{
            p->b[idx] = p->a[idx];
        }
    }
}

void tensor_abs(float *a, float *b, int size){
    int thread_num = std::thread::hardware_concurrency();
    std::vector<std::thread> threads(thread_num);
    std::vector<As> mes(thread_num);
    int blocksdim;
    if (size%thread_num != 0) {
        if (size < thread_num) {
            for (int i = 0; i < size; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = As{start, end, a, b};
            }
        }else{
            blocksdim = size/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = size;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = As{start, end, a, b};
            }
        }
    }else{
        blocksdim = size/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = As{start, end, a, b};
            }
    }
    if (size < thread_num) {
        thread_num = size;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(abs_t, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
}

// 向量扩充
typedef struct{
    int idx;
    int dim;
    int row_nums;
    int size; // 需要复制的数据的个数
    float *a; // 原数据,形状必须是向量
    float *b;
}Cot;

void copy_t(Cot *p){
    if (p->dim == 0){
        for (int i = 0; i < p->size; i++) {
            p->b[p->idx * p->size + i] = p->a[i];
        }
    }else if(p->dim == 1){
        for (int i = 0; i < p->size; i++) {
            p->b[i * p->row_nums + p->idx] = p->a[i];
        }
    }
    
}

void sscp(float *a, float *b, int a_size, int dim, int nums){
    int thread_num = std::thread::hardware_concurrency();
    std::vector<std::thread> threads(thread_num);
    std::vector<Cot> mes(thread_num);

    int blocksdim;
    if (nums%thread_num != 0) {
        if (nums < thread_num) {
            for (int i = 0; i < nums; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = Cot{i, dim, nums, a_size, a, b};
            }
        }else{
            blocksdim = nums/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = nums;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = Cot{i, dim, nums, a_size, a, b};
            }
        }
    }else{
        blocksdim = nums/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = Cot{i, dim, nums, a_size, a, b};
            }
    }
    if (nums < thread_num) {
        thread_num = nums;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(copy_t, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
}

/**************************************************/
// 矩阵求和
typedef struct{
    int start;
    int end;
    int dim;
    int m;
    int n;
    float *a;
    float *b;
}Nos;

void sum_t(Nos *p){
    
    for (int idx = p->start; idx < p->end; idx++) {
        if (p->dim == 0){
        for(int k = 0 ; k < p->m; k++){
            p->b[idx] += p->a[k * p->n + idx];
            }
        }else if(p->dim == 1){
            for(int k = 0 ; k < p->n; k++){
                p->b[idx] += p->a[idx * p->n + k];
            }
        }
    }
}

void sum(float *a, float *b, int m, int n, int dim){
    int size;
    if (dim == 0){
        size = n;
    }
    else if(dim == 1){
        size = m;
    }
    int thread_num = std::thread::hardware_concurrency();
    std::vector<std::thread> threads(thread_num);
    std::vector<Nos> mes(thread_num);
    int blocksdim;
    if (size%thread_num != 0) {
        if (size < thread_num) {
            for (int i = 0; i < size; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = Nos{start, end, dim, m, n, a, b};
            }
        }else{
            blocksdim = size/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = size;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = Nos{start, end, dim, m, n, a, b};
            }
        }
    }else{
        blocksdim = size/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = Nos{start, end, dim, m, n, a, b};
            }
    }
    if (size < thread_num) {
        thread_num = size;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(sum_t, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
}
// 均值
void mean_t(Nos *p){
    for (int idx = p->start; idx < p->end; idx++) {
        if (p->dim == 0){
        for(int k = 0 ; k < p->m; k++){
            p->b[idx] += p->a[k * p->n + idx];
            }
            p->b[idx] /= p->m;
        }else if(p->dim == 1){
            for(int k = 0 ; k < p->n; k++){
                p->b[idx] += p->a[idx * p->n + k];
            }
            p->b[idx] /= p->n;
        }
    }
}

void mean(float *a, float *b, int m, int n, int dim){
    int size;
    if (dim == 0){
        size = n;
    }
    else if(dim == 1){
        size = m;
    }
    int thread_num = std::thread::hardware_concurrency();
    std::vector<std::thread> threads(thread_num);
    std::vector<Nos> mes(thread_num);
    int blocksdim;
    if (size%thread_num != 0) {
        if (size < thread_num) {
            for (int i = 0; i < size; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = Nos{start, end, dim, m, n, a, b};
            }
        }else{
            blocksdim = size/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = size;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = Nos{start, end, dim, m, n, a, b};
            }
        }
    }else{
        blocksdim = size/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = Nos{start, end, dim, m, n, a, b};
            }
    }
    if (size < thread_num) {
        thread_num = size;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(mean_t, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
}


// 最大值
void max_t(Nos *p){
    for (int idx = p->start; idx < p->end; idx++) {
        float max;
        if (p->dim == 0){
            max = p->a[idx];
            for(int k = 0 ; k < p->m; k++){
                if(p->a[k * p->n + idx] > max){
                    max = p->a[k * p->n + idx];
                    p->b[idx] = (k * p->n + idx);
                }
            }
        }else if(p->dim == 1){
            max = p->a[idx * p->n];
            for(int k = 0 ; k < p->n; k++){
                if(p->a[idx * p->n + k] > max){
                    max = p->a[idx * p->n + k];
                }
            }
        }
    }
}
/**************************************************/
